#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
为 default_user 创建测试数据，解决 conversation_search 返回空结果的问题
"""

import asyncio
import sys
import os
from pathlib import Path
from datetime import datetime

# 添加项目根目录到路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / "src"))

# 设置环境变量
os.environ['PYTHONPATH'] = str(project_root / "src")

async def create_test_data():
    """为 default_user 创建测试数据"""
    print("=== 为 default_user 创建测试数据 ===")
    
    try:
        # 1. 导入必要模块
        print("\n1. 导入模块...")
        
        # 直接连接 ChromaDB 添加数据
        import chromadb
        from sentence_transformers import SentenceTransformer
        
        # 2. 连接 ChromaDB
        print("\n2. 连接 ChromaDB...")
        
        data_dir = project_root / "data" / "chroma_db"
        chroma_client = chromadb.PersistentClient(path=str(data_dir))
        print(f"✓ ChromaDB 连接成功: {data_dir}")
        
        # 3. 获取或创建 conversations 集合
        print("\n3. 获取 conversations 集合...")
        
        try:
            conversations_collection = chroma_client.get_collection('conversations')
            print("✓ conversations 集合已存在")
        except Exception:
            # 如果集合不存在，创建它
            conversations_collection = chroma_client.create_collection(
                name='conversations',
                metadata={"hnsw:space": "cosine"}
            )
            print("✓ 创建了新的 conversations 集合")
        
        # 4. 检查当前数据
        print("\n4. 检查当前数据...")
        current_count = conversations_collection.count()
        print(f"当前 conversations 集合文档数量: {current_count}")
        
        # 5. 准备测试数据
        print("\n5. 准备测试数据...")
        
        test_conversations = [
            {
                "id": "conv_default_user_001",
                "content": "王者荣耀是什么游戏？",
                "metadata": {
                    "user_id": "default_user",
                    "conversation_id": "conv_001",
                    "role": "user",
                    "timestamp": datetime.now().isoformat(),
                    "message_id": "msg_001"
                }
            },
            {
                "id": "conv_default_user_002",
                "content": "王者荣耀是腾讯开发的一款MOBA类手机游戏，是目前最受欢迎的手游之一。",
                "metadata": {
                    "user_id": "default_user",
                    "conversation_id": "conv_001",
                    "role": "assistant",
                    "timestamp": datetime.now().isoformat(),
                    "message_id": "msg_002"
                }
            },
            {
                "id": "conv_default_user_003",
                "content": "王者荣耀中哪个英雄最厉害？",
                "metadata": {
                    "user_id": "default_user",
                    "conversation_id": "conv_001",
                    "role": "user",
                    "timestamp": datetime.now().isoformat(),
                    "message_id": "msg_003"
                }
            },
            {
                "id": "conv_default_user_004",
                "content": "王者荣耀中不同英雄有不同的定位，比较热门的包括：李白、韩信、貂蝉、诸葛亮、孙悟空等。每个英雄都有其独特的技能和玩法。",
                "metadata": {
                    "user_id": "default_user",
                    "conversation_id": "conv_001",
                    "role": "assistant",
                    "timestamp": datetime.now().isoformat(),
                    "message_id": "msg_004"
                }
            },
            {
                "id": "conv_default_user_005",
                "content": "王者荣耀游戏有什么技巧吗？",
                "metadata": {
                    "user_id": "default_user",
                    "conversation_id": "conv_002",
                    "role": "user",
                    "timestamp": datetime.now().isoformat(),
                    "message_id": "msg_005"
                }
            },
            {
                "id": "conv_default_user_006",
                "content": "王者荣耀游戏技巧包括：合理选择英雄、掌握技能连招、注意地图意识、团队配合等。多练习是提高技术的关键。",
                "metadata": {
                    "user_id": "default_user",
                    "conversation_id": "conv_002",
                    "role": "assistant",
                    "timestamp": datetime.now().isoformat(),
                    "message_id": "msg_006"
                }
            }
        ]
        
        print(f"准备添加 {len(test_conversations)} 条测试数据")
        
        # 6. 添加测试数据到 ChromaDB
        print("\n6. 添加测试数据...")
        
        # 提取数据
        ids = [conv["id"] for conv in test_conversations]
        documents = [conv["content"] for conv in test_conversations]
        metadatas = [conv["metadata"] for conv in test_conversations]
        
        # 添加到集合
        conversations_collection.add(
            ids=ids,
            documents=documents,
            metadatas=metadatas
        )
        
        print(f"✓ 成功添加 {len(test_conversations)} 条测试数据")
        
        # 7. 验证数据添加
        print("\n7. 验证数据添加...")
        
        new_count = conversations_collection.count()
        print(f"添加后 conversations 集合文档数量: {new_count}")
        print(f"新增文档数量: {new_count - current_count}")
        
        # 8. 测试搜索功能
        print("\n8. 测试搜索功能...")
        
        # 测试目标搜索参数
        target_query = "王者荣耀"
        target_user_id = "default_user"
        
        # 无过滤搜索
        results_no_filter = conversations_collection.query(
            query_texts=[target_query],
            n_results=5
        )
        print(f"搜索 '{target_query}' (无过滤): {len(results_no_filter['ids'][0])} 个结果")
        
        # 带用户ID过滤的搜索
        results_with_filter = conversations_collection.query(
            query_texts=[target_query],
            n_results=5,
            where={"user_id": target_user_id}
        )
        print(f"搜索 '{target_query}' (用户: {target_user_id}): {len(results_with_filter['ids'][0])} 个结果")
        
        # 显示搜索结果详情
        if results_with_filter['ids'][0]:
            print("\n搜索结果详情:")
            for i, (doc_id, content, metadata, distance) in enumerate(zip(
                results_with_filter['ids'][0],
                results_with_filter['documents'][0],
                results_with_filter['metadatas'][0],
                results_with_filter['distances'][0]
            )):
                print(f"  结果 {i+1}:")
                print(f"    ID: {doc_id}")
                print(f"    内容: {content[:80]}..." if len(content) > 80 else f"    内容: {content}")
                print(f"    用户ID: {metadata.get('user_id')}")
                print(f"    相似度距离: {distance:.4f}")
        
        # 9. 验证其他搜索词
        print("\n9. 验证其他搜索词...")
        
        other_queries = ["游戏", "英雄", "技巧"]
        
        for query in other_queries:
            results = conversations_collection.query(
                query_texts=[query],
                n_results=3,
                where={"user_id": target_user_id}
            )
            print(f"搜索 '{query}' (用户: {target_user_id}): {len(results['ids'][0])} 个结果")
        
        print("\n=== 测试数据创建完成 ===")
        print("\n现在可以重新测试 MCP conversation_search 工具:")
        print(f"搜索参数: {{\"query\": \"{target_query}\", \"user_id\": \"{target_user_id}\", \"limit\": 5}}")
        print("预期结果: 应该能找到相关的对话数据")
        
    except Exception as e:
        print(f"\n创建测试数据失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(create_test_data())